#coding=utf-8

from tgrocery import Grocery
import pandas as pd
import jieba

stopwordsfiles = "../pyltp_server/stop_words.txt"
segment_dict = "../pyltp_server/pyltp_seg_dict.txt"

# 导入停顿词
with open(stopwordsfiles, 'r') as fr:
    lines = fr.readlines()
    stopwords = [word for word in lines]

# 导入分词
with open(segment_dict, 'r') as fr:
    lines = fr.readlines()
    for line in lines:
        jieba.add_word(line.strip())


def deleteStopWords(words):
    new_words = []
    for word in words:
        if not word in stopwords:
            new_words.append(word)
    return new_words

# 加载文件，导入数据,分词
def loadfile():
    meeting_data = pd.read_csv("../../data/train_data/meeting_schedule.csv")
    task_data = pd.read_csv("../../data/train_data/task_schedule.csv")
    schedule_data = pd.read_csv("../../data/train_data/schedule.csv")
    not_schedule_data = pd.read_csv("../../data/train_data/not_schedule.csv")
    test_data = pd.read_csv("../../data/train_data/schedule_testset.csv")

    def cw(contents, labels):
        result_l = []
        for content, label in zip(contents, labels):
            content = content.replace('\n', '')
            words = list(jieba.cut(content))
            new_words = []
            for w in words:
                if not w in stopwords:
                    new_words.append(w)
            result_l.append((label, ''.join(new_words)))
        return result_l

    meeting_words = cw(meeting_data['content'], ['IsMeetingSchedule'] * len(meeting_data['content']))
    taks_words = cw(task_data['content'],['TaskSchedule'] * len(task_data['content']))
    schedule_words = cw(schedule_data['content'],['IsSchedule'] * len(schedule_data['content']))
    not_schedule_words = cw(not_schedule_data['content'][:6000],['NotSchedule'] * 6000)
    test_data = cw(test_data['content'], test_data['label'])

    train_data = meeting_words + taks_words + schedule_words + not_schedule_words


    return train_data, test_data

def train():
    train_data, test_data = loadfile()
    grocery = Grocery('sample')
    print("training ......")
    grocery.train(train_data)
    grocery.save()
    print(grocery.test(test_data))


def predict(sentent):
    new_grocery = Grocery('sample')
    new_grocery.load()
    sentent = ' '.join(jieba.cut(sentent))
    pred_label = new_grocery.predict(sentent)
    return pred_label

if __name__ == '__main__':
    # train()

    sentent = '明天下午三点开会'
    print(predict(sentent))